from procinfomod import *
import time
import copy
import os, sys, getopt
import collectd

procmon = None
interval = 5
hadoopuser = 11
hdfsuser = 64
nummaps = 6
numreduces = 4

def configure(ObjConf):
  pass

def initialize():
  global procmon
  procmon = ProcMon()

def read():
  global interval
  procmon.monitor()

def sendMetric(time, name, met_name, group, val):
  global interval
  metric = collectd.Values()
  metric.time = time
  metric.plugin = group
  metric.plugin_instance = name
  metric.type = "gauge"
  metric.type_instance = met_name
  metric.interval = interval
  metric.values = [float(val)]
  metric.dispatch()
  '''
  if debug == 1:
    print name, met_name, group, val
  '''
  '''
  cmdExec = cmdGmetric + " --name="+name+ " --value="+str(val) \
            + " --group="+group + " --type="+vtype \
            + " --tmax=" + str(interval*2) \
            + " --dmax=" + str(interval*20)
  if debug == 1:
    print "cmdExec="+cmdExec
  else:
    gmetricResult = os.system(cmdExec)
  '''

def sendout(time, name, group, mets, met_name, met_st):
  for i in range(0, len(mets)):
    for j in range(met_st[i], len(mets[i])):
      val = mets[i][j]
      sendMetric(time, name, met_name[i][j], group, val)

def sendoutmr(time, name, group, mets, met_name, met_st, prefix):
  for i in range(0, len(mets)-1):
    for j in range(met_st[i], len(mets[i])):
      val = mets[i][j]
      sendMetric(time, name, prefix+"_"+met_name[i][j], group, val)

def sendoutmr_task(time, name, group, mets, met_name, met_st):
  for i in range(0, len(mets)-1):
    for j in range(met_st[i], len(mets[i])):
      val = mets[i][j]
      sendMetric(time, name, met_name[i][j], group, val)

class ProcInfo:
  def __init__(self, pid, plist, modules):
    self.pid = pid
    self.plist = plist
    self.time = 0
    self.modules = modules
    self.metrics = []
    for module in self.modules:
      self.metrics.append([0] * module.size())
    self.moreInit()

  def moreInit(self):
    self.cmd = getCmdLine(self.pid)
    (self.ppid, self.uid) = getPPidAndUid(self.pid)

  def update(self, nowtime):
    oldtime = self.time
    if nowtime - oldtime > interval * 2:
      self.moreInit()
    i = 0
    for module in self.modules:
      module.update(self.pid, oldtime, nowtime, self.metrics[i])
      i += 1
    self.time = nowtime

  def prepare_tree(self):
    self.child = []

  def create_tree(self):
    if self.ppid in self.plist.procs:
      self.parent = self.plist.procs[self.ppid]
      self.parent.child.append(self.ppid)
    else:
      self.parent = None

  def get_subtree(self):
    queue = [self.pid]
    tree = set([])    #Prevent there is a cycle
    tree.add(self.pid)
    head = 0
    while head < len(tree):
      v = queue[head]
      head += 1
      if v in self.plist.procs:
        p = self.plist.procs[v]
      else:
        continue
      for ch in p.child:
        if ch not in tree:
          tree.add(ch)
          queue.append(ch)
    return tree

  def query(self):
    return self.metrics

class JobStore:
  def __init__(self, met_name, met_st):
    self.jobs = {}
    self.met_name = met_name
    self.met_st = met_st

  def ready(self, time):
    for key, val in self.jobs.items():
      updatetime = val[-1]
#      if time - updatetime > interval * 5:
      if time > updatetime:
        del self.jobs[key]
      else:
        for i in range(0, len(val)-1):
          for j in range(0, len(val[i])):
            val[i][j] = 0

  def update(self, jobid, met, time):
    if jobid not in self.jobs:
      val = copy.deepcopy(met)
      val.append(time)
      self.jobs[jobid] = val
      return
    val = self.jobs[jobid]
    for i in range(0, len(met)):
      for j in range(0, len(met[i])):
        val[i][j] += met[i][j]
    val[-1] = time

  def output(self, time, slotmap):
    for jobid, met in self.jobs.items():
      sendoutmr(time, jobid[0:17], "mrjob", met, self.met_name, self.met_st, slotmap.mapslot(jobid))
      sendoutmr_task(time, jobid[0:28], "mrjob_task", met, self.met_name, self.met_st)

class TaskSlotMap:
  def __init__(self):
    self.taskids = [[None] * nummaps, [None] * numreduces]
    self.times = [[0] * nummaps, [0] * numreduces]
    self.waittask = []

  def update(self, tid, time):
    if tid[18] == 'm':
      index = 0
    else:
      index = 1
    i = 0
    for taskid in self.taskids[index]:
      if taskid == tid:
        self.times[index][i] = time
        return
      i += 1
    self.waittask.append(tid)

  def clean(self, updatetime):
    global interval
    for i in range(0, 2):
      j = 0
      for time in self.times[i]:
        if time < updatetime and self.taskids[i][j] is not None:
          self.taskids[i][j] = None
        j += 1

  def findslot(self, tid, time):
    if tid[18] == 'm':
      index = 0
    else:
      index = 1
    i = 0
    for taskid in self.taskids[index]:
      if taskid == None:
        self.taskids[index][i] = tid
        self.times[index][i] = time
        return
      i += 1

  def clearwait(self, time):
    for i in range(0, len(self.waittask)):
      tid = self.waittask.pop()
      self.findslot(tid, time)

  def mapslot(self, tid):
    if tid[18] == 'm':
      index = 0
    else:
      index = 1
    i = 0
    for taskid in self.taskids[index]:
      if taskid == tid:
        return tid[18]+str(i)
      i += 1
    return ''

class ProcMon:
  def __init__(self):
    # list of service processes
    self.procs = {}
    self.active = []
    self.mrjobs = []
    self.time = 0
    self.modules = [PIModIO(), PIModStatus(), PIModPS()]
    self.agg_met = []
    self.met_name = []
    self.met_st = []
    for module in self.modules:
      self.agg_met.append([0] * module.size())
      self.met_name.append(module.naming())
      self.met_st.append(module.start())
    self.jobstore = JobStore(self.met_name, self.met_st)
    self.taskslot = TaskSlotMap()

  def update(self):
    self.active = []
    self.time = int(time.time())
    for f in os.listdir("/proc/"):
      if len(f) > 0 and f[0] >= '0' and f[0] <= '9':
        pid = int(f)
        self.active.append(pid)
        if pid not in self.procs:
          self.procs[pid] = ProcInfo(pid, self, self.modules)
        self.procs[pid].update(self.time)

  def aggregate(self, active):
    for i in range(0, len(self.agg_met)):
      for j in range(self.met_st[i], len(self.agg_met[i])):
        self.agg_met[i][j] = 0
        for pid in active:
          self.agg_met[i][j] += self.procs[pid].metrics[i][j]

#Improvement: not need to scan the proc tree everytime?

  def classify(self, pinfo):
    name = ""
    group = ""
    if pinfo.uid == hadoopuser or pinfo.uid == hdfsuser or True:
      group = "hadoop"
      if pinfo.cmd.find("DataNode") >= 0:
        name = "DataNode"
      elif pinfo.cmd.find("TaskTracker") >= 0:
        name = "TaskTracker"
      elif pinfo.cmd.find("NameNode") >= 0:
        name = "NameNode"
      elif pinfo.cmd.find("JobTracker") >= 0:
        name = "JobTracker"
      else:
        pos = pinfo.cmd.find("attempt")
        if pos >= 0:
          pos = pos + len("attempt")
          name = pinfo.cmd[pos+1:pos+27]
          group = "mrjob"
    return (group, name)

  def monitor(self):
    now = int(time.time())
    self.update()
    self.mrjobs = []
    for pid in self.active:
      pinfo = self.procs[pid]
      pinfo.prepare_tree()
    for pid in self.active:
      pinfo = self.procs[pid]
      pinfo.create_tree()
    self.jobstore.ready(now)
    for pid in self.active:
      pinfo = self.procs[pid]
      (group, pname) = self.classify(pinfo)
      if group == "mrjob":
        subtree = pinfo.get_subtree()
        self.aggregate(subtree)
        for node in subtree:
          self.mrjobs.append(node)
        self.jobstore.update(pname, self.agg_met, now)
        self.taskslot.update(pname, now)
      elif group == "hadoop" and pname != "":
        sendout(now, pname, group, pinfo.metrics, self.met_name, self.met_st)
    self.aggregate(self.active)
    sendout(now, "Total", "hadoop", self.agg_met, self.met_name, self.met_st)
    self.aggregate(self.mrjobs)
    sendout(now, "TotalMRJob", "hadoop", self.agg_met, self.met_name, self.met_st)
    self.taskslot.clean(now)
    self.taskslot.clearwait(now)
    self.jobstore.output(now, self.taskslot)

  def run(self, interval):
    while (1):
      self.monitor()
      time.sleep(interval)

if __name__ != '__main__':
  collectd.register_config(configure)
  collectd.register_init(initialize)
  collectd.register_read(read, interval)

if __name__ == '__main__':
# Main Program
# init
  try:
    opts, args = getopt.getopt(sys.argv[1:], "i:d", ["interval=", "debug=", "gdir=", "log="])
  except getopt.GetoptError, err:
    print str(err)
    usage()
    sys.exit(2)

  # default value
  interval = 5
  debug = 1
  gmetricdir = "/usr/bin"
  for o, a in opts:
    if o in ("-i", "--interval"):
      interval = int(a)
    elif o in ("-d", "--debug"):
      debug = 1
    elif o in ("--gdir"):
      gmetricdir = a
    elif o in ("--log"):
      logfile = a

  cmdGmetric = gmetricdir+"/gmetric"
  if os.path.exists(cmdGmetric) == False:
    print "Cannot find gmetric in the path "+gmetricdir
    sys.exit(2)

  procmon = ProcMon()
  procmon.run(interval)
